import asyncio
from datetime import datetime
import json
import re
from mitmproxy import http, websocket
from mitmproxy.addonmanager import Loader
from surrealdb import Surreal
import logging
from addons.sse import responseheaders

from db import DB
from addons.utils.request_data_schema import RequestDataSchema
from addons.x10jqka.kaifangshi_fund import KaifangshiFund
from addons.scrape_urls import SCRAPE_URLS


class Main:
    def __init__(self):
        self.count = 0
        self.srrldb = None  # 用于存储 Surreal DB 的实例
        # why using tasks, because surrealdb python api looks like not working concurrently
        self.tasks = [] 

    def load(self, loader: Loader):
        logging.info("load()...........")

    async def running(self):
        logging.info("\n\n LLLLLLL Called when the proxy is completely up and running. At this point, you can expect all addons to be loaded and all options to be set.")
        _db = DB()
        await _db.connect()
        self.srrldb = _db.srrldb
        logging.info(f"AsyncClass loaded srrldb.client_state= {self.srrldb.client_state }")

        # Start a new asynchronous task to monitor the tasks list
        asyncio.create_task(self.monitor_tasks())

    async def monitor_tasks(self):
        wait_seconds = 0
        task_count = 0
        while True:
            # Wait for the first task to complete
            if self.tasks:
                wait_seconds = 0
                task_count = task_count + 1
                await self.tasks.pop(0) # process task one by one and await the task has done
                logging.info(f"\n*******\nmonitor_tasks \ntask_count={task_count} left tasks={len(self.tasks)}\n")
            else:
                wait_seconds = 3
            # Introduce a delay before checking the tasks again
            await asyncio.sleep(wait_seconds)  # Adjust the delay as needed

    async def done(self):
        if self.monitor_task:
            self.monitor_task.cancel()  # Cancel the monitoring task

    async def request(self, flow: http.HTTPFlow) -> None:
        # Handle HTTP/HTTPS requests as before
        pretty_url = flow.request.pretty_url
        logging.info(f">>>>>> request() pretty_url={pretty_url} \n" )
        for url_pattern_id, url_pattern in SCRAPE_URLS.items():
            if re.match(url_pattern, pretty_url):
                data = self.get_request(url_pattern_id,flow)
                if data:
                    create_task = self.srrldb.create('flow_request', data)
                    # Add the create task to the list without await
                    self.tasks.append(create_task)

    async def response(self, flow: http.HTTPFlow) -> None:
        # Handle HTTP/HTTPS requests as before
        pretty_url = flow.request.pretty_url
        logging.info(f"<<<<< response() pretty_url={pretty_url} \n" )
        for url_pattern_id, url_pattern in SCRAPE_URLS.items():
            # logging.info(f"\n\>>>> pattern= {pattern} handler= {handler} for url={url}")
            if re.match(url_pattern, pretty_url):
                data = self.get_response(url_pattern_id,flow)
                if data:
                    create_task = self.srrldb.create('flow_response', data)
                    # Add the create task to the list without await
                    self.tasks.append(create_task)

    async def websocket_message(self, flow: http.HTTPFlow):
        assert flow.websocket is not None  # make type checker happy
        pretty_url = flow.request.pretty_url
        if flow.response:
            content_type = flow.response.headers.get("Content-Type")
            if content_type:
                logging.info(f"***websocket_message()** Content-Type: {content_type}" )
            else:
                logging.info(f"***websocket_message()** Content-Type header not found in the response: {flow.response!r}")
        else:
            logging.info(f"***websocket_message()** No response available for the WebSocket message")  

        logging.info(f"<<<>>> websocket_message() pretty_url={pretty_url} \n")

        # get the latest message
        message = flow.websocket.messages[-1]

        # Check for Content-Type header
        content_type_header = message.headers.get("Content-Type")

        # was the message sent from the client or server?
        if message.from_client:
            logging.info(f"Client sent content_type_header={content_type_header} a message: {message.content!r}\n")
        else:
            logging.info(f"Server sent content_type_header={content_type_header} a message: {message.content!r}\n")

        # manipulate the message content
        message.content = re.sub(rb"^Hello", b"HAPPY", message.content)


    def get_request(self,url_pattern_id, flow) :
        pretty_url = flow.request.pretty_url

        return {
            "flow_id": flow.id,
            "url_pattern_id": url_pattern_id,
            "request_timestamp": str(datetime.utcnow()),
            "request_pretty_url": pretty_url,
            "status": "REQUESTED"
        }


    def get_response(self, url_pattern_id, flow) :
        if flow.response and len(flow.response.content) > 0:        
            # 获取响应头的内容类型值
            content_type = flow.response.headers.get("content-type")                

            # Parse the charset from the Content-Type header
            charset = "utf-8"  # Default to UTF-8 if not specified
            if content_type and "charset=" in content_type:
                charset = content_type.split("charset=")[-1].strip()
            
            # Check if the content is likely to be text-based
            if content_type and ("text/" in content_type or "application/json" in content_type):
                try:
                    # Decode the content using the specified charset
                    text_content = flow.response.content.decode(charset)
                    # Now you can work with the text_content as a string
                    return  {
                            "flow_id": flow.id,
                            "url_pattern_id": url_pattern_id,
                            "response_timestamp": str(datetime.utcnow()),
                            "response_charset": charset,
                            "response_content_type": content_type,
                            "response_content": text_content,
                            "status":"RESPONSED"
                        }
                except UnicodeDecodeError:
                    logging.info(
                        f"Failed to decode content using charset: {charset} for flow.id={flow.id}")
            else:
                logging.info(f"Non-text content, skipping decoding for flow.id={flow.id} flow.response={flow.response!r}")        

addons = [
    Main(),
    responseheaders,
]
